## code to run muni-level MSMs and other models with non-imputed data ##

msm_muni_noimp <- function(dv1, dv2, dvp, indata1, indata2, num){
  
  print(paste0("MUNI ", num))
  
  af.muni.done <- indata1 %>%
    select(GEOID,
           treat_2006,
           treat_2022,
           zri_2022_fin_st,
           zri_2006_fin_st,
           pop_total_2022,
           pop_density_2022,
           per_18t34_2022,
           per_35t64_2022,
           per_65over_2022,
           entropy_er_2022,
           per_child_2022,
           per_nomove_2022,
           hownrt_2022,
           per_immigrant_2022,
           median_yr_str_2022,
           median_hhld_inc_2022,
           hhld_pov_rt_2022,
           unem_rt_2022,
           gini_2022,
           pop_total_2009,
           pop_density_2009,
           per_18t34_2009,
           per_35t64_2009,
           per_65over_2009,
           entropy_er_2009,
           per_child_2009,
           per_nomove_2009,
           hownrt_2009,
           per_immigrant_2009,
           median_yr_str_2009,
           median_hhld_inc_2009,
           hhld_pov_rt_2009,
           unem_rt_2009,
           gini_2009,
           mgr_npov_2009,
           mgr_npov_2022,
           mgr_hpov_2009,
           mgr_hpov_2022,
           mgr_hpov_cc_2009,
           mgr_hpov_cc_2022,
           rb_npov_2009,
           rb_npov_2022,
           rb_hpov_2009,
           rb_hpov_2022,
           rb_hpov_cc_2009,
           rb_hpov_cc_2022,
           mpv_npov_2009,
           mpv_npov_2022,
           mpv_hpov_2009,
           mpv_hpov_2022,
           mpv_hpov_cc_2009,
           mpv_hpov_cc_2022,
           mgr_npov_rel_2009,
           mgr_npov_rel_2022,
           mgr_hpov_rel_2009,
           mgr_hpov_rel_2022,
           mgr_hpov_cc_rel_2009,
           mgr_hpov_cc_rel_2022,
           mpv_npov_rel_2009,
           mpv_npov_rel_2022,
           mpv_hpov_rel_2009,
           mpv_hpov_rel_2022,
           mpv_hpov_cc_rel_2009,
           mpv_hpov_cc_rel_2022)

  prop.table(table(af.muni.done$treat_2006,useNA="ifany"))

  af.muni.fe.done <- indata2
  
  prop.table(table(af.muni.fe.done$treat,useNA="ifany"))
  
  vars.2022 <- paste(c("pop_total_2022","pop_density_2022",
                       "per_18t34_2022","per_35t64_2022",
                       "per_65over_2022","entropy_er_2022",
                       "per_child_2022", "per_nomove_2022",
                       "hownrt_2022", "per_immigrant_2022",
                       "median_yr_str_2022","median_hhld_inc_2022",
                       "hhld_pov_rt_2022","unem_rt_2022",
                       "gini_2022"
  ), collapse = "+")
  
  vars.2009 <- paste(c("pop_total_2009","pop_density_2009",
                       "per_18t34_2009","per_35t64_2009",
                       "per_65over_2009", "entropy_er_2009",
                       "per_child_2009", "per_nomove_2009",
                       "hownrt_2009", "per_immigrant_2009",
                       "median_yr_str_2009","median_hhld_inc_2009",
                       "hhld_pov_rt_2009","unem_rt_2009",
                       "gini_2009"
  ), collapse = "+")
  
  fe.vars <- paste(c("pop_total","pop_density",
                     "per_18t34","per_35t64",
                     "per_65over","entropy_er",
                     "per_child", "per_nomove",
                     "hownrt", "per_immigrant",
                     "median_yr_str","median_hhld_inc",
                     "hhld_pov_rt","unem_rt",
                     "gini","GEOID","tp"
  ), collapse = "+")
  
  ## regression formulas ##
  
  form1 <- paste0("treat_2022 ~ ",
                  dv1, " + treat_2006 +",
                  vars.2022)
  
  form2 <- paste0("treat_2006 ~ ", 
                  vars.2009)
  
  form3 <- paste0(dv2," ~",
                  "thist")
  
  form4 <- paste0(dv2," ~",
                  "zri_2022_fin_st +",
                  "zri_2006_fin_st +",
                  dv1," +",
                  vars.2022," +",
                  vars.2009)
  
  form5 <- paste0(dvp," ~",
                  "zri_fin"," +",
                  fe.vars)
  
  dv1 <- ensym(dv1)
  
  muni.2022.comp <- af.muni.done %>%
    filter(!is.na(!!dv1))
  
  if(nrow(muni.2022.comp) < 30){
    
    print("TOO FEW OBS, RETURNING NA")
    return(NA)
  }

  
  muni.2022 <- glm(as.formula(form1),
                  data = muni.2022.comp,
                  binomial(link = 'logit'))
  
  ## predicted probabilities ##
  
  ppout.2022 <- predict(muni.2022, type = "response", se.fit = TRUE)
  pps.2022 <- ppout.2022$fit

  ## check pps ##
  summary(pps.2022)

  ## numerator for stabilized weights ##
  
  mean(muni.2022.comp$treat_2022)
  pps.2022.num <- rep(mean(muni.2022.comp$treat_2022),
                      length(pps.2022))
  
  ## check probabilities ##
  summary(pps.2022.num)

  ## now 2006 ##
  
  muni.2006.comp <- af.muni.done %>%
    filter(!is.na(treat_2006) & 
           GEOID %in% muni.2022.comp$GEOID)

  muni.2006 <- glm(as.formula(form2),
                  data = muni.2006.comp,
                  binomial(link = 'logit'))
  
  summary(muni.2006)
  
  ppout.2006 <- predict(muni.2006, type = "response", se.fit = TRUE)
  pps.2006 <- ppout.2006$fit

  ## check pps ##
  summary(pps.2006)

  ## numerator for stabilized weights ##

  mean(muni.2006.comp$treat_2006,na.rm=T)
  pps.2006.num <- rep(mean(muni.2006.comp$treat_2006,na.rm=T),
                      length(pps.2006))
  
  ## check numerator ##
  summary(pps.2006.num)
  
  print("LENGTHS - NOIMP")
  print(length(pps.2006))
  print(length(pps.2022))

  ## attach the pps and numerators and create weights ##
  
  af.muni.wts <- muni.2022.comp %>%
    mutate(thist = treat_2006 + treat_2022,
           num_t_2006 = pps.2006.num,
           pps_2006 = pps.2006,
           num_t_2022 = pps.2022.num,
           pps_2022 = pps.2022,
           num_2006 = ifelse(treat_2006 == 1, num_t_2006, 1-num_t_2006),
           den_2006 = ifelse(treat_2006 == 1, pps_2006, 1-pps_2006),
           num_2022 = ifelse(treat_2022 == 1, num_t_2022, 1-num_t_2022),
           den_2022 = ifelse(treat_2022 == 1, pps_2022, 1-pps_2022),
           swt_2006 = num_2006/den_2006,
           swt_2022 = num_2022/den_2022,
           swt_fin = swt_2006 * swt_2022)
  
  ## truncate the weights ##
  
  wts.cb.p5 <- quantile(af.muni.wts$swt_fin, 0.05, na.rm=T)
  wts.cb.p95 <- quantile(af.muni.wts$swt_fin,0.95, na.rm=T)
  
  af.muni.wts$swt_fin[af.muni.wts$swt_fin > wts.cb.p95] <- wts.cb.p95
  af.muni.wts$swt_fin[af.muni.wts$swt_fin < wts.cb.p5] <- wts.cb.p5
  summary(af.muni.wts$swt_fin)
  sd(af.muni.wts$swt_fin, na.rm=T)
  
  ## assign the weights for export ##
  wts.fin <- af.muni.wts$swt_fin
  
  ## msm ##
  
  muni.msm <- lm(as.formula(form3),
                data = af.muni.wts,
                weights=swt_fin)
  
  ## unadjusted ## 
  
  muni.noadj <- lm(as.formula(form3),
                  data = af.muni.wts)
  
  ## adl ## 
  
  muni.adl <- lm(as.formula(form4),
                data = af.muni.wts)
  
  ## fixed effects ## 
  
  muni.fe <- lm(as.formula(form5),
               data = af.muni.fe.done)
  
  out.msm <- summary(muni.msm)$coefficients
  out.noadj <- summary(muni.noadj)$coefficients
  out.adl <- summary(muni.adl)$coefficients
  out.fe <- summary(muni.fe)$coefficients
  
  ## assess balance ##
  
  af.muni.bal <- af.muni.wts %>%
    filter(!is.na(treat_2006))
  
  bal.t2.vars <- c("pop_total_2022","pop_density_2022",
                   "per_18t34_2009","per_35t64_2009",
                   "per_65over_2009", "entropy_er_2009",
                   "per_child_2022", "per_nomove_2022",
                   "per_immigrant_2022","hownrt_2022",
                   "median_yr_str_2022","median_hhld_inc_2022",
                   "hhld_pov_rt_2022","unem_rt_2022","gini_2022")
  
  bal.t1.vars <- c("pop_total_2009","pop_density_2009",
                   "per_18t34_2009","per_35t64_2009",
                   "per_65over_2009", "entropy_er_2009",
                   "per_child_2009", "per_nomove_2009",
                   "per_immigrant_2009","hownrt_2009",
                   "median_yr_str_2009","median_hhld_inc_2009",
                   "hhld_pov_rt_2009","unem_rt_2009","gini_2009")
  
  bal.t2 <- function(var){
    
    unwt2 <- lm(as.formula(paste0(var,
                           " ~ treat_2022 + treat_2006")),
                data = af.muni.bal)
    
    out2a <- summary(unwt2)
    
    
    wt2 <- lm(as.formula(paste0(var,
                         " ~ treat_2022 + treat_2006")),
              data = af.muni.bal,
              weights=swt_fin)
    
    out2b <- summary(wt2)
    
    outres <- list(out2a,out2b)
    
    return(outres)
    
  }
  
  t2.bal <- lapply(bal.t2.vars, bal.t2)  
  
  bal.t1 <- function(var){
    
    unwt1 <- lm(as.formula(paste0(var,
                           " ~ treat_2006")),
                data = af.muni.bal)
    
    out1a <- summary(unwt1)
    
    
    wt1 <- lm(as.formula(paste0(var,
                         " ~ treat_2006")),
              data = af.muni.bal,
              weights=swt_fin)
    
    out1b <- summary(wt1)
    
    outres <- list(out1a,out1b)
    
    return(outres)
    
  }
  
  t1.bal <- lapply(bal.t1.vars, bal.t1)  
  
  ## capture all output ## 
  
  res <- list(out.msm,
              wts.fin,
              out.noadj,
              out.adl,
              out.fe,
              t2.bal,
              t1.bal)
  
  return(res)
  
}

